package com.yeeiee.catalog;

import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import cn.hutool.db.ds.DSFactory;
import cn.hutool.db.ds.hikari.HikariDSFactory;
import cn.hutool.setting.Setting;
import com.yeeiee.config.JdbcCatalogConfig;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.ConfigOption;
import org.apache.flink.table.catalog.UniqueConstraint;
import org.apache.flink.table.catalog.exceptions.CatalogException;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.TemporaryClassLoaderContext;

import javax.annotation.Nullable;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;
import java.util.regex.Pattern;

/**
 * 实现类需要有空参构造器，不然SPI无法创建对象
 */
@Slf4j
@NoArgsConstructor
public abstract class AbstractJdbcCatalog implements ClinkCatalog {
    private ClassLoader classLoader;
    private String defaultDatabase;
    private String defaultSchema;
    private DSFactory dsFactory;
    private DataSource dataSource;
    @Getter
    private String url;
    @Getter
    private String username;
    @Getter
    private String password;

    public AbstractJdbcCatalog(
            ClassLoader classLoader,
            @NonNull String url,
            @NonNull String defaultDatabase,
            @Nullable String driverName,
            @Nullable String defaultSchema,
            @Nullable String username,
            @Nullable String password) {

        this.classLoader = classLoader;
        this.url = url;
        this.defaultDatabase = defaultDatabase;
        this.defaultSchema = defaultSchema;
        this.username = username;
        this.password = password;

        val setting = new Setting("db.setting");
        setting.set("url", url);
        setting.set("driver", driverName);
        setting.set("username", username);
        setting.set("password", password);
        dsFactory = DSFactory.setCurrentDSFactory(new HikariDSFactory(setting));
        dataSource = dsFactory.getDataSource();
    }

    protected abstract Pattern urlPattern();

    @Override
    public Set<ConfigOption<?>> requiredOptions() {
        val required = new HashSet<ConfigOption<?>>();
        required.add(JdbcCatalogConfig.url);
        required.add(JdbcCatalogConfig.username);
        required.add(JdbcCatalogConfig.password);
        return required;
    }

    @Nullable
    @Override
    public String getDefaultDatabase() {
        return defaultDatabase;
    }

    @SneakyThrows
    protected Entity queryOne(String sql, Object... params) {
        return Db.use(dataSource).queryOne(sql, params);
    }

    @SneakyThrows
    protected List<Entity> queryList(String sql, Object... params) {
        return Db.use(dataSource).query(sql, params);
    }

    @SneakyThrows
    protected void execute(String sql, Object... params) {
        Db.use(dataSource).execute(sql, params);
    }


    @Override
    public void open() throws CatalogException {
        // load the Driver use userClassLoader explicitly, see FLINK-15635 for more detail
        try (val ignored = TemporaryClassLoaderContext.of(classLoader)) {
            // test connection, fail early if we cannot connect to database
            try (val conn = dataSource.getConnection()) {
                log.info("Catalog {} established connection to {}", factoryIdentifier(), url);
            } catch (SQLException e) {
                throw new CatalogException(String.format("Failed connecting to %s via JDBC.", url), e);
            }
        }
    }

    @Override
    public void close() throws CatalogException {
        dsFactory.close();
    }

    protected Optional<UniqueConstraint> getPrimaryKey(String database, String schema, String table) {
        try (val connection = dataSource.getConnection()) {
            val metaData = connection.getMetaData();
            val resultSet = metaData.getPrimaryKeys(database, schema, table);

            val keySeqColumnName = new HashMap<Integer, String>();
            String pkName = null;
            while (resultSet.next()) {
                val columnName = resultSet.getString("COLUMN_NAME");
                pkName = resultSet.getString("PK_NAME"); // all the PK_NAME should be the same
                val keySeq = resultSet.getInt("KEY_SEQ");
                Preconditions.checkState(!keySeqColumnName.containsKey(keySeq - 1), "The field(s) of primary key must be from the same table.");
                keySeqColumnName.put(keySeq - 1, columnName); // KEY_SEQ is 1-based index
            }
            val pkFields = Arrays.asList(new String[keySeqColumnName.size()]); // initialize size
            keySeqColumnName.forEach(pkFields::set);
            if (!pkFields.isEmpty()) {
                // PK_NAME maybe null according to the javadoc, generate an unique name in that case
                pkName = pkName == null ? "pk_" + String.join("_", pkFields) : pkName;
                return Optional.of(UniqueConstraint.primaryKey(pkName, pkFields));
            }
            return Optional.empty();
        } catch (SQLException e) {
            throw new CatalogException("get primary key field from: " + database + "." + table);
        }
    }
}
